#ifndef BTL_DEMO_RANSACDEMO_QRANSACDEMOAPPLICATION
#define BTL_DEMO_RANSACDEMO_QRANSACDEMOAPPLICATION

#include <btl/extra/Gui/QImageView.hpp>
#include <btl/extra/Gui/ImagePainter.hpp>

#include <btl/RobustEstimation/BasicRANSAC.hpp>
#include <btl/RobustEstimation/PreemptiveRANSAC.hpp>

#include <Eigen/Core>
#include <Eigen/SVD>
#include <boost/random.hpp>

#include <QKeyEvent>
#include <QApplication>
#include <QTimer>

#include <iterator>
#include <iostream>
#include <fstream>

inline Eigen::Vector2d least_squares_line(const std::vector<Eigen::Vector2d>& obs)
{
   Eigen::MatrixXd A(obs.size(), 2);
   Eigen::VectorXd b(obs.size());

   A.col(1).setOnes();

   for (int i = 0; i < obs.size(); ++i)
   {
      A(i,0) = obs[i](0);
      b(i) = obs[i](1);
   }

   // find least-squares solution for x in Ax = b
   return A.jacobiSvd(Eigen::ComputeThinU | Eigen::ComputeThinV).solve(b);
}

inline double score_hypothesis(const Eigen::Vector2d& line, const Eigen::Vector2d& point)
{
   // distance function
   const double m = line(0);
   const double c = line(1);

   double dist = m*point(0) + c - point(1);
   return dist*dist;
}

class LineFitSearchSpace : public btl::ModelSearchSpace<Eigen::Vector2d>
{
   public:
      LineFitSearchSpace() {}

      std::vector<Eigen::Vector2d> points;

      virtual int numObservations() const
      {
         return points.size();
      }

      virtual Eigen::Vector2d buildModel(const int observations[], int numObservations) const
      {
         std::vector<Eigen::Vector2d> sample;
         sample.reserve(numObservations);
         for (int i = 0; i < numObservations; ++i)
            sample.push_back(points[observations[i]]);

         return least_squares_line(sample);
      }

      virtual double errorMetric(const Eigen::Vector2d& model, int observation) const
      {
         return score_hypothesis(model, points[observation]);
      }
};

class QRANSACDemoWindow : public btl::extra::QImageViewRGB
{
      Q_OBJECT
   public:
      QRANSACDemoWindow(QWidget* parent = 0): btl::extra::QImageViewRGB(parent)
      {
         rng.seed((unsigned int)time(0));

         _pRANSAC = new btl::BasicRANSAC<Eigen::Vector2d, 2>(0.2, 10, 50, 0);
         //_pRANSAC = new btl::PreemptiveRANSAC<Eigen::Vector2d,2>(10, 0.2, 10);

         //read in data points
         std::ifstream pointsFile;
         pointsFile.open("/home/pgrad/glynn/shared/btl2/points.txt");
         Eigen::Vector2d point;
         pointsFile >> point(0) >> point(1);
         do
         {
            fitter.points.push_back(point);
            pointsFile >> point(0) >> point(1);
            //std::cout << point.transpose() << "\n";
         } while (pointsFile.good());
         pointsFile.close();

         //std::cout << "Read " << fitter.points.size() << " points\n";

         runRANSAC();
      }

   protected:
      void keyPressEvent(QKeyEvent* evt)
      {
         if (evt->key() == Qt::Key_Space)
            runRANSAC();
         else
            btl::extra::QImageViewRGB::keyPressEvent(evt);
      }

   private:
      void runRANSAC() {
         boost::random_number_generator<boost::mt19937> std_rng(rng);

         const int N = fitter.points.size();

         boost::optional<Eigen::Vector2d> solution
            = _pRANSAC->findModel(fitter, std_rng);

         if (! solution)
            std::cout << "RANSAC couldn't find a solution :/" << std::endl;

         //display results
         btl::PixelRGB white(255, 255, 255);
         btl::PixelRGB black(0, 0, 0);
         btl::PixelRGB red(255, 0, 0);
         btl::PixelRGB green(0, 255, 0);
         btl::PixelRGB blue(0, 0, 255);
         btl::ImageRGB plot(IMAGE_WIDTH, IMAGE_HEIGHT, white);
         btl::extra::ImagePainter<btl::PixelRGB> painter(plot);
         painter.drawLine(GRAPH_BORDER, IMAGE_HEIGHT-GRAPH_BORDER, GRAPH_BORDER, GRAPH_BORDER, black);
         painter.drawLine(GRAPH_BORDER, IMAGE_HEIGHT-GRAPH_BORDER, IMAGE_WIDTH-GRAPH_BORDER, IMAGE_HEIGHT-GRAPH_BORDER, black);
         for(int i = 0; i < N; ++i)
         {
            const Eigen::Vector2d pt = fitter.points[i];
            Eigen::Vector2i pixel = convertPointToPixel(pt(0), pt(1));
            //std::cout << "Pixel(" << pixel(0) << "," << pixel(1) << std::endl;
            painter.drawCircle(pixel, 1,  blue);
         }

         if (solution)
         {
            const std::vector<int>& inliers = _pRANSAC->getInliers();
            for(int i = 0; i < int(inliers.size()); ++i)
            {
               const Eigen::Vector2d pt = fitter.points[inliers[i]];
               Eigen::Vector2i pixel = convertPointToPixel(pt(0), pt(1));
               //std::cout << "Pixel(" << pixel(0) << "," << pixel(1) << std::endl;
               painter.drawCircle(pixel, 3,  green);
            }

            double y0, y1;
            Eigen::Vector2i pt0, pt1;

            //std::cout << "Solution: " << (*solution) << "\n";
            // draw solution found by RANSAC

            y0 = (*solution)(1); // c
            y1 = (*solution)(0) + y0; // m + c
            pt0 = convertPointToPixel(0.0, y0);
            pt1 = convertPointToPixel(1.0, y1);
            painter.drawLine(pt0, pt1, green);
         }

         this->setImage(plot);
      }

      Eigen::Vector2i convertPointToPixel(double x, double y)
      {
         static const int GRAPH_WIDTH = IMAGE_WIDTH - 2 * GRAPH_BORDER;
         static const int GRAPH_HEIGHT = (IMAGE_HEIGHT - 2 * GRAPH_BORDER);
         Eigen::Vector2d pixel(
            GRAPH_WIDTH * x + GRAPH_BORDER,
            IMAGE_HEIGHT - (GRAPH_HEIGHT * (y / 6.0)  + GRAPH_BORDER)
         );
         return pixel.cast<int>();
      }

      LineFitSearchSpace fitter;
      btl::BasicRANSAC<Eigen::Vector2d, 2>* _pRANSAC;
      //btl::RobustModelFinder<Eigen::Vector2d>* _pRANSAC;

      boost::mt19937 rng;

      static const int IMAGE_WIDTH = 640;

      static const int IMAGE_HEIGHT = 480;

      static const int GRAPH_BORDER = 10;
};

class QRANSACDemoApplication : public QApplication
{

      Q_OBJECT

   public:

      QRANSACDemoApplication(int& argc, char** argv) : QApplication(argc, argv)
      {
         _pDemo = new QRANSACDemoWindow;
         _pDemo->show();
      }

      ~QRANSACDemoApplication() {
         delete _pDemo;
      }

   protected:

      QRANSACDemoWindow* _pDemo;

};

#endif //BTL_DEMO_RANSACDEMO_QRANSACDEMOAPPLICATION
